# -*- coding:utf-8 -*-
"""
@file name  : lesson-05-autograd
@author     : QuZhang
@date       : 2020-11-26
@brief      : 自动求导
"""
import torch
torch.manual_seed(10)  # 设置固定的随机数种子，使得每次运行该文件时生成的随机数相同

# 一、torch.autograd.backward
# 一、backward的参数
# backward用于求解所有参数的梯度，底层调用torch.autograd.backward

# ------------- retain_graph--------------
# 功能：报存计算图
# flag = True
flag = False
if flag:
    # 1.创建张量
    w = torch.tensor([1.], requires_grad=True)
    x = torch.tensor([2.], requires_grad=True)

    # 2.前向传播，创建计算图
    a = torch.add(w, x)
    b = torch.add(w, 1)
    y = torch.mul(a, b)

    # 3.通过计算图反向传播求梯度
    # y.backward()  # 使用对象的方法来计算导数
    # print(w.grad)
    # 注：当执行完反向传播之后，默认会把计算图删除，无法再次反向求导
    # y.backward()

    # 设置报存计算图，可以实现多次反向传播
    y.backward(retain_graph=True)
    print(w.grad)
    y.backward()
    print(x.grad)


# ------------- grad_tensors ------------
# 功能：多个梯度时，为每个梯度设置相应的权重，用于求解总梯度
# flag = True
# flag = False
if flag:
    w = torch.tensor([1.], requires_grad=True)
    x = torch.tensor([2.], requires_grad=True)

    a = torch.add(w, x)
    b = torch.add(w, 1)

    y0 = torch.mul(a, b)  # y0 = (x+w) * (w+1)
    y1 = torch.add(a, b)  # y1 = (x+w) + (w+1)

    loss = torch.cat([y0, y1], dim=0)  # 拼接成[y0, y1]
    # print(loss)
    grad_tensors = torch.tensor([1., 2.])  # 梯度的权重,y0对应的权重为1.0,y1对应的梯度为2.0

    # 对每一个变量的梯度，为梯度权重*原梯度的和
    loss.backward(gradient=grad_tensors)
    print(w.grad)

# 二、torch.autograd.grad
# --------------- create_graph -----------
# 功能：创建导数的计算图,用于求高阶导
# flag = True
flag = False
if flag:
    x = torch.tensor([3.], requires_grad=True)
    y = torch.pow(x, 2)

    # y为要求导的张量，x为梯度张量
    grad_1 = torch.autograd.grad(y, x, create_graph=True)  # grad_1 = dy / dx = 2x = 6
    print(grad_1)

    # 二阶导
    grad_2 = torch.autograd.grad(grad_1[0], x)  # grad_2 = d(dy/dx) / dx = 2
    print(grad_2)

# ----------- tips: 1 ------------
# 多次反向传播时，要注意将上一次计算的梯度清零
# flag = True
flag = False
if flag:
    # 1.创建张量
    w = torch.tensor([1.], requires_grad=True)
    x = torch.tensor([2.], requires_grad=True)

    for i in range(3):
        # 2.前向传播，创建计算图
        a = torch.add(w, x)
        b = torch.add(w, 1)
        y = torch.mul(a, b)

        y.backward()
        print(w.grad)

        # 梯度清零，避免计算梯度时发生叠加
        w.grad.zero_()